import numpy as np
import torch  # 导入PyTorch，用于加载保存的模型
import torch.nn as nn  # 导入nn模块，用于定义神经网络模型
from PIL import Image
import gradio as gr
from modelstruct import MLP  # 导入自定义的多层感知机模型

model=MLP()

# 使用模型参数载入模型
params_path = 'mnist_mlp.pth'  # 定义保存模型参数的文件路径
loaded_params = torch.load(params_path)  # 使用torch.load函数从文件中载入模型参数
loaded_model_params = MLP()  # 实例化一个新的感知机模型，用于加载模型参数
loaded_model_params.load_state_dict(loaded_params)  # 使用load_state_dict函数将载入的参数加载到模型中
loaded_model_params.eval()  # 将模型设置为评估模式，关闭梯度计算和批归一化等操作

# 定义预处理函数
def preprocess(image):
    image = Image.fromarray(image)               # 将数组转换为图像对象
    image = image.resize((28, 28)).convert('L')  # 调整图像大小并转换为灰度图像
    image_array = np.array(image)                # 将图像转换为numpy数组
    normalized_image = image_array / 255.0       # 将像素值归一化到[0,1]范围内
    return normalized_image

# 定义预测函数，这个函数将用于Gradio接口进行预测
def predict(image):
    preprocessed_image = preprocess(image)                                      # 预处理输入的图像
    preprocessed_image = torch.tensor(preprocessed_image, dtype=torch.float32)  # 转换为PyTorch张量
    preprocessed_image = preprocessed_image.unsqueeze(0)                        # 增加一维，符合模型输入要求
    with torch.no_grad():  # 关闭梯度计算
        output = loaded_model_params(preprocessed_image)                        # 使用加载的模型进行预测
    predicted_digit = torch.argmax(output).item()                               # 获取预测结果
    return str(predicted_digit)

# 创建Gradio接口，这个接口将用于用户输入和显示预测结果
iface = gr.Interface(fn=predict, inputs='sketchpad', outputs='label')

# 启动Gradio接口，用户可以通过这个接口进行交互
iface.launch()